
from collections import defaultdict
from typing import Dict, Any

import networkx as nx
from python.lib.typing import Tuple, List

from allocate_switches_with_es_optimized import allocate_electric_with_es_optimized


# 计算不需要重构通信对链路利用率
def calculate_no_recon_utilization(
        G: nx.Graph,
        no_recon_pairs: list,
        current_paths:dict,
        traffic_demand: dict,
        tor_switches: set,
        k_paths: int = 30
) -> dict:
    """
    自动计算不需要重构通信对的链路利用率

    参数:
        G: 网络拓扑
        no_recon_pairs: 不需要重构的通信对列表
        traffic_demand: 流量需求字典
        k_paths: 候选路径数

    返回:
        { (u,v): used_gbps } 链路利用率字典
    """
    print(f"输出no_recon_pairs:{no_recon_pairs}")
    print(type(current_paths))
    print(f"输出current_path:")
    print(current_paths)
    link_util = defaultdict(float) # 初始化链路利用率为0


    for pair in no_recon_pairs:
        s, d = pair
        if pair in traffic_demand:
            demand_gb = traffic_demand[pair]
            # print(f"  正在更新电网络路由: {s} → {d},流量需求为{demand_gb}")
        if pair in current_paths:
            # print(f"通信对{pair}在current_paths中")
            # print(f"通信对{pair}在current_paths中,路径为{current_paths[pair]}")
            path= current_paths[pair]
            # 计算当前的链路利用率
            # valid = True
            temp_util = link_util.copy() # 创建临时变量：避免修改原始数据
            # 输出link_util
            # print(f"  当前link_util为: {temp_util}")
            for u, v in zip(path, path[1:]):
                edge = tuple(sorted((u, v)))
                capacity = G[u][v].get('capacity', 0)
                # print(f"  链路 {edge} 容量为{capacity}")
                # if (capacity - temp_util[edge]) < demand_gb:
                #     print(f"  链路 {edge} 容量不足")
                #     valid = False
                #     break
                temp_util[edge] += demand_gb
            # if valid:
        # 更新实际利用率
        for u, v in zip(path, path[1:]):
            edge = tuple(sorted((u, v)))
            link_util[edge] += demand_gb
            # print(f"  更新链路 {edge} 已使用容量为{link_util[edge]}")
        # break

    return dict(link_util)
# 为需要重构的链路找到最优的电力路径
def find_optimal_electric_paths(
        G: nx.Graph,
        recon_pairs: list,
        traffic_demand: dict,
        no_recon_pairs: list,  # 新增不需要重构的通信对列表
        existing_util: dict,  # 其他已有利用率
        tor_switches: set,
        k_paths: int = 30
) -> tuple:
    """
    全自动链路分配主函数

    返回:
        (recon_paths, total_utilization)
        - recon_paths: { (s,d): (path, delay) }
        - total_utilization: { (u,v): total_used_gbps }
    """
    # ===== 步骤1：计算不需要重构的链路利用率 =====
    no_recon_util = calculate_no_recon_utilization(
        G, no_recon_pairs, traffic_demand, tor_switches, k_paths
    )
    for edge, used in no_recon_util.items():
        pass
        # print(f"  {edge}: 已使用容量为{used}")

    # ===== 步骤2：初始化总利用率 =====
    total_util = defaultdict(float)
    for u, v in G.edges():
        edge = tuple(sorted((u, v)))
        total_util[edge] = existing_util.get(edge, 0.0) + no_recon_util.get(edge, 0.0)
        # print(f"在给需要重构的通信对分配路径之前初始化已使用的链路容量为{edge}: {total_util[edge]}")

    # ===== 步骤3：处理需要重构的通信对 =====
    recon_paths = {}
    sorted_recon = sorted(recon_pairs, key=lambda x: traffic_demand.get(x, 0), reverse=True)

    for pair in sorted_recon:
        s, d = pair
        demand_gb = traffic_demand.get(pair, 0)
        # print(f"正在更新电网络路由: {s} → {d},流量需求为{demand_gb}")

        try:
            # 获取候选路径（仅电链路）
            candidate_paths = []
            for path in nx.shortest_simple_paths(G, s, d, weight='hops'):
                if len(candidate_paths) >= k_paths:
                    break
                if all(G[u][v].get('type', 'electrical') == 'electrical'
                       for u, v in zip(path, path[1:])):
                    candidate_paths.append(path)
            for path in candidate_paths:
                pass
                # print(f"  {pair}的候选路径为: {'→'.join(path)}")

            # 评估路径可行性
            feasible = []
            for path in candidate_paths:
                temp_util = total_util.copy()
                # print(f"  当前total_util为: {temp_util}")
                valid = True
                delay = 0.0

                for u, v in zip(path, path[1:]):
                    edge = tuple(sorted((u, v)))
                    capacity = G[u][v]['capacity']
                    # print(f"  链路 {edge} 容量为{capacity}")
                    remaining = capacity - temp_util[edge]
                    # print(f"  链路 {edge} 剩余容量为{remaining}")

                    if remaining < demand_gb:
                        valid = False
                        break

                    # 时延计算（考虑传播延迟）
                    tx_time = demand_gb / remaining  # 秒
                    # prop_delay = G[u][v].get('delay', 1e-6)  # 默认为1μs
                    delay += tx_time
                    # print(f"  链路 {edge} 时延为{delay}")
                    temp_util[edge] += demand_gb
                    # print(f"  更新链路 {edge} 已使用容量为{temp_util[edge]}")

                if valid:
                    feasible.append((path, delay))
                    # print(f"  可行路径: {'→'.join(path)}")

            # 选择最优路径
            if feasible:
                best_path, min_delay = min(feasible, key=lambda x: x[1])
                # print(f"  最优路径: {'→'.join(best_path)}")
                recon_paths[pair] = (best_path, min_delay)
                # print(f"  最小时延为: {min_delay:.2f} s")

                # 更新总利用率
                for u, v in zip(best_path, best_path[1:]):
                    edge = tuple(sorted((u, v)))
                    total_util[edge] += demand_gb
                    # print(f"  更新链路 {edge} 已使用容量为{total_util[edge]}")

        except nx.NetworkXNoPath:
            continue

    return recon_paths, dict(total_util)

# 定义重构延迟，如果在电链路上的传播时间大于重构时间，则可以考虑在重构完成之后继续在最优的光链路上传播
def find_recon_electric_optical_paths(
        G: nx.Graph,
        # recon_pairs: list,
        traffic_demand: dict,
        # no_recon_pairs: list,  # 新增不需要重构的通信对列表
        # existing_util: dict,  # 其他已有利用率
        current_paths:dict,
        recon_electric_delay_info:dict,
        recon_delay: float,
        # tor_switches: set,
        # k_paths: int = 5
) -> dict[Any, Any]:
    recon_optical_delays={}
    for pair in recon_electric_delay_info:
        print("==================测试测试1======================")
        print(recon_electric_delay_info[pair])
        print("==================测试测试======================")
        # print(f"  {pair}的电链路传播时间为{recon_electric_delay_info[pair]}")
        if recon_electric_delay_info[pair] > recon_delay:
            print(f"  {pair}的电链路传播时间为{recon_electric_delay_info[pair]}大于重构时间，考虑在重构完成之后继续在最优的光链路上传播")
            # 找到最优的光链路
            path = current_paths[pair]
            print(f"通信对{pair}的重构之后最优路径为{path}")
            # 找到需要在光链路上传输的剩余流量:流量矩阵需求-电链路的传输速度*重构时间
            print(f"  通信对{pair}的初始流量需求为{traffic_demand[pair]}")
            print(f"在电链路上传输的量为{traffic_demand[pair]/recon_electric_delay_info[pair]*recon_delay}")
            remaining_traffic = traffic_demand[pair] - traffic_demand[pair]/recon_electric_delay_info[pair]*recon_delay
            # 剩余流量在已知光链路上的传输时间
            remaining_delay = remaining_traffic*8 / G[path[0]][path[1]]['capacity']
            print(f"  剩余流量{remaining_traffic}在已知光链路上的传输时间: {remaining_delay:.2f} s")
            # 将剩余流量在已知光链路上的传输时间加入到recon_optical_delays中
            recon_optical_delays[pair] = remaining_delay
        else:
            # recon_optical_delays[pair] = 0.0
            print(f"  {pair}的电链路传播时间为{recon_electric_delay_info[pair]}小于重构时间，不需要在重构完成之后继续在最优的光链路上传播")

    return recon_optical_delays

def initialize_utilization(G: nx.Graph) -> Dict[Tuple[str, str], float]:
    """初始化链路利用率字典，键为排序后的节点元组"""
    return {
        tuple(sorted((u, v))): 0.0
        for u, v in G.edges()
    }


def calculate_communication_metrics(
        topology: nx.Graph,
        traffic_matrix: Dict[Tuple[str, str], float],
        communication_pairs: List[Tuple[str, str]],
        paths: Dict[Tuple[str, str], List[str]]
) -> Tuple[Dict[Tuple[str, str], float], Dict[Tuple[str, str], float]]:
    """
    计算每个通信对的通信时间及链路流量分配

    参数:
        topology: 网络拓扑图，边需包含 'bandwidth' (Gbps) 和 'delay' (秒) 属性
        traffic_matrix: 流量需求矩阵，键为通信对，值为流量 (Gbit)
        communication_pairs: 需要分析的通信对列表
        paths: 每个通信对的预设路径，键为通信对，值为节点列表

    返回:
        Tuple[
            Dict[通信对, 通信时间(秒)],
            Dict[链路, 总流量(Gbit)]
        ]
    """
    # 初始化链路流量分配字典（统一处理无向图）
    link_allocations = initialize_utilization(topology)
    # print(f"初始化链路流量分配字典: ")
    for edge in link_allocations:
        pass
        # print(f"  初始化链路 {edge} 的流量分配为{link_allocations[edge]}")

    # 初始化通信时间字典
    # communication_times: Dict[Tuple[str, str], float] = {}
    # for pair in communication_pairs:
    #     print(f"初始化通信时间为{communication_times[pair]}")
    communication_times = {}

    for pair in communication_pairs:
        # 获取流量需求
        traffic = traffic_matrix.get(pair, 0.0)
        # print(f"获取通信对 {pair} 的流量需求为{traffic}")
        if traffic <= 0:
            continue

        # 获取路径
        path = paths.get(pair, [])
        if len(path) < 2:
            print(f"警告: 通信对 {pair} 的路径无效")
            continue


        for i in range(len(path) - 1):
            u, v = path[i], path[i + 1]
            edge = tuple(sorted((u, v)))
            # if edge in link_allocations:
            #     link_allocations[edge] += traffic
            # else:
            #     print(f"警告: 路径 {path} 包含不存在的链路 {edge}")

        # ========== 计算通信时间 ==========
        total_time = 0.0
        for i in range(len(path) - 1):
            u, v = path[i], path[i + 1]
            edge = tuple(sorted((u, v)))

            # 获取链路属性
            edge_data = topology.edges.get(edge, {})
            bandwidth = edge_data.get('capacity', 0.0)  # Gbps
            # print(f"  获取链路 {edge} 的带宽为{bandwidth}Gbps")
            # delay = edge_data.get('delay', 0.0)  # 秒

            # 计算传输时间（单位：秒）
            if bandwidth > 0:
                # print(f"链路{(u, v)}已经分配的流量为{link_allocations[edge]}")
                transmission_time = (link_allocations[edge] + traffic*8) / bandwidth
                # print(f"  更新链路 {edge} 的传输时间为{transmission_time}")
            else:
                transmission_time = float('inf')
            # ========== 更新链路流量分配 ==========
            link_allocations[edge] += traffic
            # print(f"  更新链路 {edge} 的流量分配为{link_allocations[edge]}")

            # 累加时间
            total_time += transmission_time
            # print(f"  更新通信对 {pair} 的通信时间为{total_time}")

        communication_times[pair] = total_time

    return communication_times, link_allocations


def calculate_network_metrics(
        G: nx.Graph,
        traffic_matrix: Dict[Tuple[str, str], float],  # 单位：Gbps
        communication_pairs: List[Tuple[str, str]],
        paths: Dict[Tuple[str, str], List[str]]
) -> Tuple[Dict[Tuple[str, str], float], Dict[Tuple[str, str], float]]:
    """
    计算网络通信指标

    返回:
        (通信对时间字典, 链路流量字典)
        时间单位：秒
        流量单位：Gbps
    """
    # 初始化数据结构
    communication_times = {}
    link_allocations = {tuple(sorted(e)): 0.0 for e in G.edges()}

    for pair in communication_pairs:
        src, dst = pair
        if src not in G or dst not in G:
            print(f"无效通信对: {pair}，节点不存在")
            continue

        traffic = traffic_matrix.get(pair, 0.0)
        if traffic <= 0:
            print(f"忽略零流量通信对: {pair}")
            continue

        path = paths.get(pair, [])
        if len(path) < 2:
            print(f"无效路径: {pair} -> {path}")
            continue

        # 计算路径总时间（传输时间 + 传播延迟）
        total_time = 0.0
        for i in range(len(path) - 1):
            u, v = path[i], path[i + 1]
            edge = tuple(sorted((u, v)))

            # 获取链路属性
            edge_data = G.edges.get(edge, {})
            bandwidth = edge_data.get('capacity', 0.0)  # Gbps
            # delay = edge_data.get('delay', 0.0)  # 秒

            # 计算传输时间（数据量/带宽）
            if bandwidth > 0:
                transmission_time = traffic*8 / bandwidth
            else:
                transmission_time = float('inf')
                print(f"链路 {edge} 带宽为0，无法传输")

            total_time += transmission_time

            # 累加链路流量
            if edge in link_allocations:
                link_allocations[edge] += traffic
            else:
                print(f"路径包含不存在链路: {edge}")

        communication_times[pair] = total_time

    return communication_times, link_allocations


def calculate_communication_times(
        topology: nx.Graph,
        traffic_matrix: Dict[Tuple[str, str], float],  # 单位：Gbps
        communication_pairs: List[Tuple[str, str]],
        paths: Dict[Tuple[str, str], List[str]]
) -> Dict[Tuple[str, str], float]:
    """
    计算每对通信对的总通信时间（传输时间 + 传播延迟）

    参数:
        topology: 网络拓扑图，边需包含 'bandwidth' (Gbps) 和 'delay' (秒) 属性
        traffic_matrix: 流量需求矩阵，键为通信对，值为流量需求（Gbps）
        communication_pairs: 需要分析的通信对列表
        paths: 每个通信对的预设路径，键为通信对，值为节点列表

    返回:
        字典：{通信对: 总时间（秒）}
    """
    comm_times = {}
    # 初始化每条链路已分配流量为0
    link_allocations = {tuple(sorted(e)): 0.0 for e in topology.edges()}

    for pair in communication_pairs:
        src, dst = pair

        # 1. 验证通信对有效性
        if src not in topology or dst not in topology:
            print(f"警告：无效通信对 {pair}，节点不存在")
            comm_times[pair] = float('inf')
            continue

        # 2. 获取流量需求
        traffic = traffic_matrix.get(pair, 0.0)
        # print(f"获取通信对 {pair} 的流量需求为 {traffic} Gbps")
        if traffic <= 0:
            print(f"忽略零流量通信对：{pair}")
            comm_times[pair] = 0.0
            continue

        # 3. 获取路径
        path = paths.get(pair, [])
        # print(f"获取通信对 {pair} 的路径为 {path}")
        if len(path) < 2:
            print(f"路径无效：{pair} -> {path}")
            comm_times[pair] = float('inf')
            continue

        # 4. 计算路径总时间
        total_time = 0.0
        valid_path = True

        for i in range(len(path) - 1):
            u, v = path[i], path[i + 1]
            edge = tuple(sorted((u, v)))  # 统一无向边标识

            # 4.1 检查链路是否存在和是否不是电链路
            if not topology.has_edge(u, v):
                print(f"路径包含不存在链路：{edge}")
                valid_path = False
                break

            # 4.2 获取链路属性
            edge_data = topology.edges[edge]
            bandwidth = edge_data.get('capacity', 0.0)  # Gbps
            # print(f"获取链路 {edge} 的带宽为 {bandwidth} Gbps")
            # delay = edge_data.get('delay', 0.0)  # 秒

            # 4.3 计算传输时间
            if bandwidth <= 0:
                print(f"链路 {edge} 带宽无效：{bandwidth} Gbps")
                valid_path = False
                break
            transmission_time = (link_allocations[edge]+traffic*8) / bandwidth  # 秒
            # print(f"计算链路 {edge} 的传输时间为 （{link_allocations[edge]}+{traffic}）*8 /{bandwidth}= {transmission_time} 秒")
            total_time += transmission_time
            # print(f"总时间：{total_time} 秒")
            # 更新该链路上的流量分配
            link_allocations[edge] += traffic

        if not valid_path:
            comm_times[pair] = float('inf')
            continue

        comm_times[pair] = total_time
        # print(f"通信对 {pair} 的总通信时间为 {comm_times[pair]} 秒")

    return comm_times


